#!/usr/bin/env python
"""
Celery Worker启动脚本
用于启动和管理Celery工作进程
"""
import os
import sys
import signal
import logging
from pathlib import Path

# 添加项目路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.celery_app import app
from src.utils.logging_config import setup_logging

# 设置日志
logger = setup_logging('celery_worker')

class WorkerManager:
    """Worker管理器"""
    
    def __init__(self):
        self.worker = None
        self.setup_signal_handlers()
        
    def setup_signal_handlers(self):
        """设置信号处理器"""
        signal.signal(signal.SIGTERM, self.handle_shutdown)
        signal.signal(signal.SIGINT, self.handle_shutdown)
        
    def handle_shutdown(self, signum, frame):
        """优雅关闭"""
        logger.info(f"Received signal {signum}, shutting down worker...")
        if self.worker:
            self.worker.stop()
        sys.exit(0)
        
    def start_worker(self, queues=None, concurrency=None, loglevel='INFO'):
        """
        启动Worker
        
        Args:
            queues: 要监听的队列列表
            concurrency: 并发数
            loglevel: 日志级别
        """
        worker_name = os.getenv('CELERY_WORKER_NAME', 'worker@%h')
        
        # 配置Worker参数
        worker_args = [
            'worker',
            '-n', worker_name,
            '--loglevel', loglevel,
            '--without-gossip',  # 禁用gossip以提高性能
            '--without-mingle',  # 禁用mingle以加快启动
            '--without-heartbeat',  # 禁用heartbeat以减少网络开销
            '-Ofair',  # 公平调度
        ]
        
        # 设置队列
        if queues:
            worker_args.extend(['-Q', ','.join(queues)])
        else:
            worker_args.extend(['-Q', 'high_priority,normal_priority,low_priority,crawler_queue,scheduler_queue,health_queue'])
            
        # 设置并发数
        if concurrency:
            worker_args.extend(['-c', str(concurrency)])
            
        # 启动worker
        logger.info(f"Starting Celery worker with args: {worker_args}")
        app.worker_main(argv=worker_args)

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Celery Worker Manager')
    parser.add_argument('--queues', nargs='+', help='Queues to listen')
    parser.add_argument('--concurrency', type=int, help='Worker concurrency')
    parser.add_argument('--loglevel', default='INFO', help='Log level')
    
    args = parser.parse_args()
    
    manager = WorkerManager()
    manager.start_worker(
        queues=args.queues,
        concurrency=args.concurrency,
        loglevel=args.loglevel
    )

if __name__ == '__main__':
    main()